{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#先引入后面可能用到的包（package）\n",
    "import pandas as pd  \n",
    "import numpy as np\n",
    "import talib as ta\n",
    "from datetime import datetime,timedelta\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline   \n",
    "#正常显示画图时出现的中文和负号\n",
    "from pylab import mpl\n",
    "mpl.rcParams['font.sans-serif']=['SimHei']\n",
    "mpl.rcParams['axes.unicode_minus']=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index={'上证综指': 'sh.000001',\n",
    "        '深证成指': 'sz.399001',\n",
    "        '沪深300': 'sh.000300',\n",
    "        '创业板指': 'sz.399006',\n",
    "        '上证50': 'sh.000016',\n",
    "        '中证500': 'sh.000905',\n",
    "        '中小板指': 'sz.399005',\n",
    "        '上证180': 'sh.000010'}\n",
    "#获取当前交易的股票代码和名称\n",
    "def get_code():\n",
    "    import baostock as bs\n",
    "    bs.login()\n",
    "    rs = bs.query_stock_basic()\n",
    "    bs.logout()\n",
    "    data_list = []\n",
    "    while (rs.error_code == '0') & rs.next():\n",
    "        # 获取一条记录，将记录合并在一起\n",
    "        data_list.append(rs.get_row_data())\n",
    "    df = pd.DataFrame(data_list, columns=rs.fields)\n",
    "    codes=df.code.values\n",
    "    names=df.code_name.values\n",
    "    stock=dict(zip(names,codes))\n",
    "    #合并指数和个股成一个字典\n",
    "    stocks=dict(stock,**index)\n",
    "    return stocks    \n",
    "#获取行情数据\n",
    "def get_daily_data(stock,start,end):\n",
    "    \"\"\"\n",
    "    获取历史K线数据\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    code = get_code()[stock]\n",
    "    import baostock as bs\n",
    "    bs.login()\n",
    "    rs = bs.query_history_k_data_plus(code,\"date,code,open,high,low,close,preclose,volume,amount,pctChg\",start,end,frequency='d')\n",
    "    data_list = []\n",
    "    while (rs.error_code == '0') & rs.next():\n",
    "        # 获取一条记录，将记录合并在一起\n",
    "        data_list.append(rs.get_row_data())\n",
    "    result = pd.DataFrame(data_list, columns=rs.fields)\n",
    "    result.open = result.open.astype(float)\n",
    "    result.high = result.high.astype(float)\n",
    "    result.low = result.low.astype(float)\n",
    "    result.close = result.close.astype(float)\n",
    "    result.date= pd.to_datetime(result.date)\n",
    "    result.set_index('date',inplace=True)\n",
    "    bs.logout()\n",
    "    result=result.sort_index()\n",
    "    #计算收益率\n",
    "    result['ret']=result.close/result.close.shift(1)-1\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hs=get_daily_data('沪深300','2018-01-01','2024-07-01')[['close','open','high','low','volume']]\n",
    "#最近N1个交易日最高价\n",
    "hs['up']=ta.MAX(hs.high,timeperiod=20).shift(1)\n",
    "#最近N2个交易日最低价\n",
    "hs['down']=ta.MIN(hs.low,timeperiod=10).shift(1)\n",
    "#每日真实波动幅度\n",
    "hs['ATR']=ta.ATR(hs.high,hs.low,hs.close,timeperiod=20)\n",
    "hs.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_strategy(data):\n",
    "    x1=data.close>data.up\n",
    "    x2=data.close.shift(1)<data.up.shift(1)\n",
    "    x=x1&x2\n",
    "    y1=data.close<data.down\n",
    "    y2=data.close.shift(1)>data.down.shift(1)\n",
    "    y=y1&y2\n",
    "    data.loc[x,'signal']='buy'\n",
    "    data.loc[y,'signal']='sell'\n",
    "    buy_date=(data[data.signal=='buy'].index).strftime('%Y%m%d')\n",
    "    sell_date=(data[data.signal=='sell'].index).strftime('%Y%m%d')\n",
    "    buy_close=data[data.signal=='buy'].close.round(2).tolist()\n",
    "    sell_close=data[data.signal=='sell'].close.round(2).tolist()\n",
    "    return (buy_date,buy_close,sell_date,sell_close)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#对K线图和唐奇安通道进行可视化\n",
    "from pyecharts.charts import Kline, Line\n",
    "from pyecharts import options as opts\n",
    "from pyecharts.globals import ThemeType\n",
    "\n",
    "# 假设hs是您的历史数据DataFrame\n",
    "attr = [str(t) for t in hs.index.strftime('%Y%m%d')]\n",
    "v1 = np.array(hs.loc[:, ['open', 'close', 'low', 'high']])\n",
    "\n",
    "# 假设构建K线图\n",
    "kline = (\n",
    "    Kline()\n",
    "    .add_xaxis(attr)\n",
    "    .add_yaxis(\"K线\", v1.tolist(), itemstyle_opts=opts.ItemStyleOpts(color=\"#ef232a\", color0=\"#14b143\"))\n",
    ")\n",
    "\n",
    "# 假设构建唐奇安通道图，这里需要根据实际数据计算上边界、下边界\n",
    "upper_bound, lower_bound = ...  # 计算唐奇安通道的上界和下界数据\n",
    "donchian_channel = (\n",
    "    Line()\n",
    "    .add_xaxis(attr)\n",
    "    .add_yaxis(\"上边界\", upper_bound.tolist(), linestyle_opts=opts.LineStyleOpts(opacity=0.5))\n",
    "    .add_yaxis(\"下边界\", lower_bound.tolist(), linestyle_opts=opts.LineStyleOpts(opacity=0.5))\n",
    ")\n",
    "\n",
    "# 使用 Page 组合多个图表\n",
    "page = Page(layout=ThemeType.LIGHT)\n",
    "page.add(kline)\n",
    "page.add(donchian_channel)\n",
    "\n",
    "# 渲染到本地网页\n",
    "page.render_notebook()  # 如果在Jupyter notebook中使用\n",
    "# 或者 page.render(\"your_chart.html\") 保存为HTML文件\n",
    "page = Page(layout=ThemeType.LIGHT)\n",
    "\n",
    "attr=[str(t) for t in hs.index.strftime('%Y%m%d')]\n",
    "v1=np.array(hs.loc[:,['open','close','low','high']])\n",
    "v2=np.array(hs.up)\n",
    "v3=np.array(hs.down)\n",
    "kline = Kline(\"沪深300唐奇安通道\",title_text_size=15)\n",
    "kline.add(\"K线图\", attr, v1.round(1),is_datazoom_show=True,)\n",
    "# 成交量\n",
    "bar = Bar()\n",
    "bar.add(\"成交量\", attr, hs['vol'],tooltip_tragger=\"axis\", is_legend_show=False, \n",
    "        is_yaxis_show=False, yaxis_max=5*max(hs[\"vol\"]))\n",
    "line = Line()\n",
    "line.add(\"上轨线\", attr, v2.round(1),is_datazoom_show=True,\n",
    "         is_smooth=True,is_symbol_show=False,line_width=1.5)\n",
    "line.add(\"下轨线\", attr, v3.round(1),is_datazoom_show=True,\n",
    "         is_smooth=True,is_symbol_show=False,line_width=1.5)\n",
    "#添加买卖信号\n",
    "bd,bc,sd,sc=my_strategy(hs)\n",
    "es = EffectScatter(\"buy\")\n",
    "es.add( \"sell\", sd, sc, )\n",
    "es.add(\"buy\", bd, bc,symbol=\"triangle\",)\n",
    "overlap = Overlap(width=2000, height=600)\n",
    "overlap.add(kline)\n",
    "overlap.add(line)\n",
    "overlap.add(bar,yaxis_index=1, is_add_yaxis=True)\n",
    "overlap.add(es)\n",
    "page.add(overlap, grid_right=\"10%\")\n",
    "page.render_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#关掉pandas的warnings\n",
    "pd.options.mode.chained_assignment = None\n",
    "def strategy(stock,start,end,N1=20,N2=10):\n",
    "    df=get_daily_data(stock,start,end)\n",
    "    #最近N1个交易日最高价\n",
    "    df['H_N1']=ta.MAX(df.high,timeperiod=N1)\n",
    "    #最近N2个交易日最低价\n",
    "    df['L_N2']=ta.MIN(df.low,timeperiod=N2)\n",
    "    #当日收盘价>昨天最近N1个交易日最高点时发出信号设置为1\n",
    "    buy_index=df[df.close>df['H_N1'].shift(1)].index\n",
    "    df.loc[buy_index,'收盘信号']=1\n",
    "    #将当日收盘价<昨天最近N2个交易日的最低点时收盘信号设置为0\n",
    "    sell_index=df[df.close<df['L_N2'].shift(1)].index\n",
    "    df.loc[sell_index,'收盘信号']=0\n",
    "    df['当天仓位']=df['收盘信号'].shift(1)\n",
    "    df['当天仓位'].fillna(method='ffill',inplace=True)\n",
    "    d=df[df['当天仓位']==1].index[0]-timedelta(days=1)\n",
    "    df1=df.loc[d:].copy()\n",
    "    df1['ret'][0]=0\n",
    "    df1['当天仓位'][0]=0\n",
    "    #当仓位为1时，买入持仓，当仓位为0时，空仓，计算资金净值\n",
    "    df1['策略净值']=(df1.ret.values*df1['当天仓位'].values+1.0).cumprod()\n",
    "    df1['指数净值']=(df1.ret.values+1.0).cumprod()\n",
    "    df1['策略收益率']=df1['策略净值']/df1['策略净值'].shift(1)-1\n",
    "    df1['指数收益率']=df1.ret\n",
    "    total_ret=df1[['策略净值','指数净值']].iloc[-1]-1\n",
    "    annual_ret=pow(1+total_ret,250/len(df1))-1\n",
    "    dd=(df1[['策略净值','指数净值']].cummax()-df1[['策略净值','指数净值']])/df1[['策略净值','指数净值']].cummax()\n",
    "    d=dd.max()\n",
    "    beta=df1[['策略收益率','指数收益率']].cov().iat[0,1]/df1['指数收益率'].var()\n",
    "    alpha=(annual_ret['策略净值']-annual_ret['指数净值']*beta)\n",
    "    exReturn=df1['策略收益率']-0.03/250\n",
    "    sharper_atio=np.sqrt(len(exReturn))*exReturn.mean()/exReturn.std()\n",
    "    TA1=round(total_ret['策略净值']*100,2)\n",
    "    TA2=round(total_ret['指数净值']*100,2)\n",
    "    AR1=round(annual_ret['策略净值']*100,2)\n",
    "    AR2=round(annual_ret['指数净值']*100,2)\n",
    "    MD1=round(d['策略净值']*100,2)\n",
    "    MD2=round(d['指数净值']*100,2)\n",
    "    S=round(sharper_atio,2)\n",
    "    df1[['策略净值','指数净值']].plot(figsize=(15,7))\n",
    "    plt.title('海龟交易策略简单回测',size=15)\n",
    "    bbox = dict(boxstyle=\"round\", fc=\"w\", ec=\"0.5\", alpha=0.9)\n",
    "    plt.text(df1.index[int(len(df1)/5)], df1['指数净值'].max()/1.5, f'累计收益率：\\\n",
    "策略{TA1}%，指数{TA2}%;\\n年化收益率：策略{AR1}%，指数{AR2}%；\\n最大回撤：  策略{MD1}%，指数{MD2}%;\\n\\\n",
    "策略alpha: {round(alpha,2)}，策略beta：{round(beta,2)}; \\n夏普比率：  {S}',size=13,bbox=bbox)  \n",
    "    plt.xlabel('')\n",
    "    ax=plt.gca()\n",
    "    ax.spines['right'].set_color('none')\n",
    "    ax.spines['top'].set_color('none')\n",
    "    plt.show()\n",
    "    #return df1.loc[:,['close','ret','H_N1','L_N2','当天仓位','策略净值','指数净值']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy('上证综指','2005-01-01','2024-07-01')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy('沪深300','','')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy('创业板指','','')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy('沪深300','2018-01-01','')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy('中国平安','2005-01-01','',N1=20,N2=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy('东方通信','2013-01-01','',N1=20,N2=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy('贵州茅台','2005-01-01','',N1=20,N2=10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "FAL",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
